from __future__ import print_function
from tensorflow import keras
from keras.datasets import cifar100
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D

import matplotlib.pyplot as plt

# 下面的代码可以加载CIFAR-100数据集
# The data, shuffled and split between train and test sets:
(x_train, y_train), (x_test, y_test) = cifar100.load_data()
print('x_train shape:', x_train.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')

# Each image is a 32 x 32 x 3 的字符数组，由于在计算机中，每个图片的像素点颜色都是用0-255范围的数字表示的
# 32 x 32 x 3 数组表示 长度为32，高度为32，深度为3，也就是有三个通道，
# RGB就是记录了黄色，绿色，蓝色的三个通道的记录方法中的一种
print(x_train[444].shape)
print(x_train[444])
# 展示一下训练前的图片
print(y_train[442])
plt.imshow(x_train[442])
# 数据集分类为100个
num_classes = 100
# 将整型的类别标签转为onehot编码，onehot编码是一种方便计算机处理的二元编码。
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)
# 展示下100长度的onehot编码，格式为（000000……10000）0占99个，1占一个
print(y_train[444])
# 归一化，转化为浮点数
x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255

# 自己构建一个模型
model_1 = Sequential()

# 模型的构建：可以理解为卷积层是提取特征，然后maxPooling是平滑他们的共性
# 5x5 convolution with 2x2 stride and 32 filters
model_1.add(Conv2D(32, (5, 5), strides=(2, 2), padding='same',
                   input_shape=x_train.shape[1:]))
model_1.add(Activation('relu'))

# Another 5x5 convolution with 2x2 stride and 32 filters
model_1.add(Conv2D(32, (5, 5), strides=(2, 2)))
model_1.add(Activation('relu'))

# 2x2 max pooling reduces to 3 x 3 x 32
model_1.add(MaxPooling2D(pool_size=(2, 2)))
model_1.add(Dropout(0.25))

# Flatten turns 3x3x32 into 288x1
model_1.add(Flatten())
model_1.add(Dense(512))
model_1.add(Activation('relu'))
model_1.add(Dropout(0.5))
model_1.add(Dense(num_classes))
model_1.add(Activation('softmax'))

model_1.summary()

# 构建并训练你的初始的卷积神经网络，以识别CIFAR-100数据集
# Let's build a CNN using Keras' Sequential capabilities
# 创建CNN模型
model_2 = Sequential()

model_2.add(Conv2D(32, (3, 3), padding='same',
                   input_shape=x_train.shape[1:]))
model_2.add(keras.layers.BatchNormalization())
model_2.add(Dropout(0.25))

model_2.add(Conv2D(32, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(MaxPooling2D(pool_size=(2, 2)))
# model_2.add(Dropout(0.25))

model_2.add(Conv2D(64, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(Dropout(0.3))

model_2.add(Conv2D(64, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(MaxPooling2D(pool_size=(2, 2)))

model_2.add(Conv2D(128, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(Dropout(0.4))

model_2.add(Conv2D(128, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(Dropout(0.4))

model_2.add(Conv2D(128, (3, 3), padding='same'))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(MaxPooling2D(pool_size=(2, 2)))

model_2.add(Flatten())
model_2.add(Dense(512))
model_2.add(Activation('relu'))
model_2.add(keras.layers.BatchNormalization())
model_2.add(Dropout(0.5))

model_2.add(Dense(num_classes))
model_2.add(Activation('softmax'))

# Check number of parameters
model_2.summary()

# 设置训练一次的批次块大小，块越大，一般训练会快，但是块太大，内存装不下也不好
batch_size = 128

# initiate RMSprop optimizer
# opt_2 = keras.optimizers.SGD(lr=0.1,decay = 1e-6,momentum=0.9, nesterov=True)
# opt = keras.optimizers.RMSprop(learning_rate=0.001, decay=1e-6)
opt = keras.optimizers.RMSprop(learning_rate=0.001)

# Let's train the model using RMSprop
model_2.compile(loss='categorical_crossentropy',
                optimizer=opt,
                metrics=['accuracy'])

history = model_2.fit(x_train, y_train,
                      batch_size=batch_size,
                      # epochs=50,
                      epochs=3,
                      validation_data=(x_test, y_test),
                      shuffle=True)

# 绘制训练 & 验证的准确率值
plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title('Model accuracy')
plt.ylabel('Accuracy')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.show()

# 绘制训练 & 验证的损失值
plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('Model loss')
plt.ylabel('Loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Test'], loc='upper left')
plt.show()
